package com.ruoyi.chat.utils;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.FileInputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * Langchain4j工具类
 *
 * @author ruoyi
 */
@Component
public class ChatLangchainUtils {
    private static final Logger log = LoggerFactory.getLogger(ChatLangchainUtils.class);

    // Redis连接信息
    //@Value("${redis.host:192.168.6.100}")
    @Value("${redis.host:47.104.144.178}")
    private String redisHost;

    @Value("${redis.port:6379}")
    private Integer redisPort;

    @Value("${spring.data.redis.password:}")
    private String redisPassword;

    // OpenAI配置
    @Value("${langchain4j.open-ai.chat-model.api-key:}")
    private String openAiApiKey;

    @Value("${langchain4j.open-ai.chat-model.base-url:https://api.openai.com/v1}")
    private String openAiBaseUrl;

    @Value("${langchain4j.open-ai.chat-model.model-name:gpt-3.5-turbo}")
    private String openAiModelName;

    @Value("${langchain4j.open-ai.chat-model.timeout:60}")
    private Integer openAiTimeout;
    // 阿里云百炼
    @Value("${langchain4j.open-ai.embedding-model.api-key:}")
    private String embeddingApiKey;

    @Value("${langchain4j.open-ai.embedding-model.base-url:}")
    private String embeddingBaseUrl;

    @Value("${langchain4j.open-ai.embedding-model.model-name:}")
    private String embeddingModelName;

    @Value("${langchain4j.open-ai.embedding-model.timeout:60}")
    private Integer embeddingTimeout;

    @Value("${langchain4j.open-ai.embedding-model.dimensions:1536}")
    private Integer embeddingDimensions;


    /**
     * 获取阿里云的Embedding模型
     *
     * @return EmbeddingModel 阿里云的Embedding模型
     */
    public EmbeddingModel getEmbeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .apiKey(embeddingApiKey)
                .baseUrl(embeddingBaseUrl)
                .modelName(embeddingModelName)
                .timeout(java.time.Duration.ofSeconds(embeddingTimeout))
                .build();
    }


    /**
     * 获取Redis向量存储
     *
     * @param namespace 命名空间，用于区分不同的知识库
     * @return EmbeddingStore Redis向量存储
     */
    public EmbeddingStore<TextSegment> getLocalEmbeddedStore(String namespace) {
        try {
            log.info("创建Redis向量存储，主机：{}，端口：{}，维度：{}", redisHost, redisPort, embeddingDimensions);
            return RedisEmbeddingStore
                    .builder()
                    .host(redisHost)
                    .port(redisPort)
                    .dimension(embeddingDimensions)
                    .build();
        } catch (Exception e) {
            log.error("创建Redis向量存储失败", e);
            // 如果Redis配置有问题，回退到内存存储，避免应用启动失败
            log.warn("回退到内存向量存储");
            return new InMemoryEmbeddingStore<>();
        }
    }


    /**
     * 获取文档分割器
     *
     * @param maxSegmentSize 最大片段大小
     * @param maxOverlap     最大重叠大小
     * @return DocumentSplitter 文档分割器
     */
    public DocumentSplitter getDocumentSplitter(int maxSegmentSize, int maxOverlap) {
        return DocumentSplitters.recursive(maxSegmentSize, maxOverlap);
    }

    /**
     * 将文件解析并向量化存储，同时返回切片数据
     *
     * @param file      文件对象
     * @param fileType  文件类型
     * @param namespace 命名空间
     * @param fileId    文件ID，用于关联切片数据
     * @return 切片数据和向量信息的包装对象
     */
    public FileSegmentResult ingestFileWithSegments(File file, String fileType, String namespace, Long fileId) {
        try {
            // 初始化进度信息
            initializeProgress(fileId, file.getName());

            log.info("开始解析文件: {}, 类型: {}, 大小: {}字节, 文件ID: {}", file.getName(), fileType, file.length(), fileId);
            updateProgress(fileId, 0, 0, "解析文件中");

            // 解析文档
            long parseStartTime = System.currentTimeMillis();
            Document document = parseDocument(file, fileType);
            long parseEndTime = System.currentTimeMillis();
            log.info("文档解析完成, 耗时: {}ms", (parseEndTime - parseStartTime));

            updateProgress(fileId, 0, 0, "文档解析完成");

            if (document == null) {
                log.error("文档解析结果为空");
                failProgress(fileId, "文档解析结果为空");
                return new FileSegmentResult(0, null);
            }

            // 判断文档是否为空
            String text = document.text();
            if (text == null || text.trim().isEmpty()) {
                log.error("文档内容为空");
                failProgress(fileId, "文档内容为空");
                return new FileSegmentResult(0, null);
            }

            log.info("文档内容提取成功, 文本长度: {} 字符", text.length());
            // 根据文件名判断是否根据5个#分割
            // 使用Langchain4j提供的文档分割器
            DocumentSplitter splitter;
            List<TextSegment> textSegments = new ArrayList<>();
            if (file.getName().contains("5#")) {
                log.info("根据5个#分割");
                // 正则表达式分割
                // joinDelimter表示将正则匹配到的部分替换成什么
                DocumentByRegexSplitter documentByRegexSplitter = new DocumentByRegexSplitter("#{5}", "", 2048, 100);
                String[] splits = documentByRegexSplitter.split(document.text());
                log.info("正则匹配到 {} 个段落", splits.length);
                // 我明白了，在通过#####分割时，会产生空白段落，而空白段不能用在TextSegment.textSegment(split);，会报错。
                for (String split : splits) {
                    // 判断是否是空白段落
                    if (split == null || split.trim().isEmpty()) {
                        continue;
                    }
                    System.out.println("=======================");
                    TextSegment segment = TextSegment.textSegment(split);
                    textSegments.add(segment);
                }
                //log.info("textsegments: {}", textSegments);
                //log.info("文档分割完成, 正则匹配到 {} 个段落", textSegments.size());
            } else {
                // 否则，使用常规分割
                log.info("常规文档分割");
                splitter = getDocumentSplitter(500, 200);
                textSegments = splitter.split(document);
            }

            // 分割文档
            log.info("开始分割文档");
            updateProgress(fileId, 0, 0, "分割文档中");

            // 过滤掉空段落
            textSegments = textSegments.stream()
                    .filter(segment -> segment.text() != null && !segment.text().trim().isEmpty())
                    .collect(Collectors.toList());
            log.info("文档分割完成, 共 {} 个有效段落", textSegments.size());
            updateProgress(fileId, 0, textSegments.size(), "文档分割完成");

            // 获取嵌入模型和存储
            // 初始化Embedding模型
            EmbeddingModel embeddingModel = getEmbeddingModel();
            log.info("Embedding模型初始化完成：{}", embeddingModel);
            // 更新向量化进度
            updateProgress(fileId, 0, textSegments.size(), "初始化向量模型完成");

            log.info("初始化向量存储, 命名空间: {}", namespace);
            EmbeddingStore<TextSegment> embeddingStore = getLocalEmbeddedStore(namespace);
            log.info("向量存储初始化完成：{}", embeddingStore);
            updateProgress(fileId, 0, textSegments.size(), "初始化向量存储完成");

            int successCount = 0;
            int failCount = 0;
            int totalCount = textSegments.size();
            int processedCount = 0;

            log.info("开始向量化处理, 共 {} 个段落需要处理", totalCount);
            updateProgress(fileId, 0, totalCount, "开始向量化处理");
            long totalStartTime = System.currentTimeMillis();

            // 保存切片数据
            List<FileSegmentInfo> segmentInfoList = new ArrayList<>();

            // 向量化并存储
            // 获取文件名，拼接到每一个切片中
            String fileName = file.getName();
            log.info("文件名称：{}", fileName);
            for (TextSegment textSegment : textSegments) {
                processedCount++;
                // 每处理10个段落记录一次进度
                if (processedCount % 10 == 0 || processedCount == totalCount) {
                    log.info("向量化进度: {}/{} ({}%)",
                            processedCount, totalCount,
                            Math.round((processedCount * 100.0) / totalCount));
                    // 更新进度信息
                    updateProgress(fileId, processedCount, totalCount, "向量化处理中");
                }
                // 在最后拼接文件名(方便后期查找文件来源)
                //String segmentText = textSegment.text().trim() + " | 以上内容来源于：" + fileName;
                TextSegment segment = TextSegment.from(textSegment.text().trim() + " | 以上内容来源于：" + fileName);
                try {
                    // 直接向量化并存储
                    Embedding embeded = embeddingModel.embed(segment).content();
                    embeddingStore.add(embeded, segment);
                    String vectorId = "vector:" + fileId + ":" + successCount;
                    FileSegmentInfo segmentInfo = new FileSegmentInfo(
                            fileId,
                            segment.text(),
                            successCount,
                            vectorId
                    );
                    segmentInfoList.add(segmentInfo);
                    successCount++;
                } catch (Exception e) {
                    failCount++;
                    log.warn("处理段落失败: {}", e.getMessage());
                    log.debug("失败段落长度: {}, 段落前100个字符: {}",
                            segment.toString().length(),
                            segment.toString().length() > 100 ? segment.toString().substring(0, 100) + "..." :
                                    segment.toString());
                }
            }
            long totalEndTime = System.currentTimeMillis();
            log.info("向量化处理完成, 总耗时: {}ms, 成功: {}, 失败: {}, 切片数据: {}",
                    (totalEndTime - totalStartTime), successCount, failCount, segmentInfoList.size());

            // 标记处理完成
            completeProgress(fileId);

            return new FileSegmentResult(successCount, segmentInfoList);
        } catch (Exception e) {
            // 标记处理失败
            failProgress(fileId, e.getMessage());

            log.error("文件向量化存储失败: {}", e.getMessage(), e);
            throw new RuntimeException("文件向量化失败: " + e.getMessage(), e);
        }
    }

    /**
     * 切片信息数据类
     */
    public static class FileSegmentInfo {
        private Long fileId;
        private String content;
        private Integer segmentIndex;
        private String vectorId;

        public FileSegmentInfo(Long fileId, String content, Integer segmentIndex, String vectorId) {
            this.fileId = fileId;
            this.content = content;
            this.segmentIndex = segmentIndex;
            this.vectorId = vectorId;
        }

        public Long getFileId() {
            return fileId;
        }

        public String getContent() {
            return content;
        }

        public Integer getSegmentIndex() {
            return segmentIndex;
        }

        public String getVectorId() {
            return vectorId;
        }
    }

    /**
     * 文件向量化结果包装类
     */
    public static class FileSegmentResult {
        private int vectorCount;
        private List<FileSegmentInfo> segments;

        public FileSegmentResult(int vectorCount, List<FileSegmentInfo> segments) {
            this.vectorCount = vectorCount;
            this.segments = segments;
        }

        public int getVectorCount() {
            return vectorCount;
        }

        public List<FileSegmentInfo> getSegments() {
            return segments;
        }
    }

    /**
     * @param file     文件对象
     * @param fileType 文件类型
     * @return Document 解析后的文档
     */
    private Document parseDocument(File file, String fileType) {
        // 使用try with resource 自动关闭流，释放系统资源
        try (FileInputStream is = new FileInputStream(file)) {
            log.info("文件类型: {}", fileType);
            ApacheTikaDocumentParser apacheTikaDocumentParser = new ApacheTikaDocumentParser();
            ApachePdfBoxDocumentParser apachePdfBoxDocumentParser = new ApachePdfBoxDocumentParser();
            TextDocumentParser textDocumentParser = new TextDocumentParser();
            // 创建文档Metadata
            Metadata metadata = new Metadata();
            metadata.put("file_name", file.getName());
            metadata.put("file_path", file.getPath());
            String filePath = String.valueOf(file);
            log.info("解析文件路径：{}", filePath);
            log.info("准备解析文档");
            // 如果是pdf类型，则使用pdfbox
            // 解析pdf文档，Tika在这里会卡死，于是使用pdfbox
            Document document;
            if ("pdf".equals(fileType)) {
                log.info("使用pdfbox解析pdf文档");
                document = apachePdfBoxDocumentParser.parse(is);
            } else if ("txt".equals(fileType)) {
                log.info("使用textDocumentParser解析txt文档");
                document = textDocumentParser.parse(is);
            } else {
                log.info("使用tika解析文档");
                document = apacheTikaDocumentParser.parse(is);
            }

            log.info("文档解析完成");
            //log.info("document: {}", document);
            // 不再根据文件类型选择不同的解析器，而是全部使用Tika
            // 判断文件当中是否有文字内容
            if (document.text() == null || document.text().trim().isEmpty()) {
                log.error("文件内容为空或解析失败: {}", file.getName());
                return null;
            }
            //log.info("解析后的文档内容：{}", document.text());
            return Document.from(document.text(), metadata);
        } catch (Exception e) {
            log.error("解析文档失败: {} - {}", e.getClass().getName(), e.getMessage(), e);
            throw new RuntimeException("解析文档失败: " + e.getMessage(), e);
        }
    }

    /**
     * 将Document对象按照Langchain4j规则进行分段处理
     *
     * @param document       要分段的文档
     * @param maxSegmentSize 最大分段大小（字符数）
     * @param overlap        重叠大小（字符数）
     * @return 分段后的TextSegment列表
     */
    public List<dev.langchain4j.data.segment.TextSegment> splitDocument(
            dev.langchain4j.data.document.Document document,
            int maxSegmentSize,
            int overlap) {

        try {
            log.info("开始分割文档，最大段落大小: {}, 重叠大小: {}", maxSegmentSize, overlap);

            // 创建文档分割器
            dev.langchain4j.data.document.DocumentSplitter splitter =
                    dev.langchain4j.data.document.splitter.DocumentSplitters.recursive(maxSegmentSize, overlap);

            // 分割文档
            List<dev.langchain4j.data.segment.TextSegment> segments = splitter.split(document);

            log.info("文档分割完成，共生成 {} 个文本段落", segments.size());
            return segments;
        } catch (Exception e) {
            log.error("文档分割处理失败: {}", e.getMessage(), e);
            throw new RuntimeException("文档分割处理失败: " + e.getMessage(), e);
        }
    }

    /**
     * 向量化进度信息类
     */
    public static class VectorizeProgressInfo {
        private Long fileId;             // 文件ID
        private String fileName;         // 文件名
        private int totalSegments;       // 总段落数
        private int processedSegments;   // 已处理段落数
        private double progress;         // 进度百分比
        private String status;           // 状态：处理中/已完成/失败
        private String message;          // 其他信息或错误消息
        private long startTime;          // 开始时间
        private long lastUpdateTime;     // 最后更新时间

        public VectorizeProgressInfo(Long fileId, String fileName) {
            this.fileId = fileId;
            this.fileName = fileName;
            this.totalSegments = 0;
            this.processedSegments = 0;
            this.progress = 0.0;
            this.status = "初始化";
            this.startTime = System.currentTimeMillis();
            this.lastUpdateTime = this.startTime;
        }

        public void update(int processedSegments, int totalSegments, String status) {
            this.processedSegments = processedSegments;
            this.totalSegments = totalSegments;
            this.progress = totalSegments > 0 ? (double) processedSegments / totalSegments * 100 : 0;
            this.status = status;
            this.lastUpdateTime = System.currentTimeMillis();
        }

        public void setError(String errorMessage) {
            this.status = "失败";
            this.message = errorMessage;
            this.lastUpdateTime = System.currentTimeMillis();
        }

        public void setComplete() {
            this.processedSegments = this.totalSegments;
            this.progress = 100.0;
            this.status = "已完成";
            this.lastUpdateTime = System.currentTimeMillis();
        }

        // Getters
        public Long getFileId() {
            return fileId;
        }

        public String getFileName() {
            return fileName;
        }

        public int getTotalSegments() {
            return totalSegments;
        }

        public int getProcessedSegments() {
            return processedSegments;
        }

        public double getProgress() {
            return progress;
        }

        public String getStatus() {
            return status;
        }

        public String getMessage() {
            return message;
        }

        public long getStartTime() {
            return startTime;
        }

        public long getLastUpdateTime() {
            return lastUpdateTime;
        }

        public long getElapsedTimeMs() {
            return lastUpdateTime - startTime;
        }
    }

    // 用于存储进度信息的Map，键为文件ID
    private static final Map<Long, VectorizeProgressInfo> progressMap = new ConcurrentHashMap<>();

    /**
     * 初始化文件处理进度
     *
     * @param fileId   文件ID
     * @param fileName 文件名
     * @return 进度信息对象
     */
    public VectorizeProgressInfo initializeProgress(Long fileId, String fileName) {
        VectorizeProgressInfo progress = new VectorizeProgressInfo(fileId, fileName);
        progressMap.put(fileId, progress);
        log.info("初始化文件 {} (ID: {}) 的向量化进度跟踪", fileName, fileId);
        return progress;
    }

    /**
     * 更新文件处理进度
     *
     * @param fileId            文件ID
     * @param processedSegments 已处理段落数
     * @param totalSegments     总段落数
     * @param status            当前状态
     */
    public void updateProgress(Long fileId, int processedSegments, int totalSegments, String status) {
        VectorizeProgressInfo progress = progressMap.get(fileId);
        if (progress != null) {
            progress.update(processedSegments, totalSegments, status);
            log.debug("更新文件 {} 的向量化进度: {}/{} ({}%), 状态: {}",
                    progress.getFileName(), processedSegments, totalSegments,
                    String.format("%.2f", progress.getProgress()), status);
        }
    }

    /**
     * 标记文件处理完成
     *
     * @param fileId 文件ID
     */
    public void completeProgress(Long fileId) {
        VectorizeProgressInfo progress = progressMap.get(fileId);
        if (progress != null) {
            progress.setComplete();
            log.info("文件 {} 向量化处理完成, 总耗时: {}ms",
                    progress.getFileName(), progress.getElapsedTimeMs());
        }
    }

    /**
     * 标记文件处理失败
     *
     * @param fileId       文件ID
     * @param errorMessage 错误信息
     */
    public void failProgress(Long fileId, String errorMessage) {
        VectorizeProgressInfo progress = progressMap.get(fileId);
        if (progress != null) {
            progress.setError(errorMessage);
            log.error("文件 {} 向量化处理失败: {}", progress.getFileName(), errorMessage);
        }
    }

    /**
     * 获取文件处理进度
     *
     * @param fileId 文件ID
     * @return 进度信息对象，如果不存在则返回null
     */
    public VectorizeProgressInfo getProgress(Long fileId) {
        return progressMap.get(fileId);
    }

    /**
     * 删除文件处理进度信息
     *
     * @param fileId 文件ID
     * @return 被删除的进度信息对象，如果不存在则返回null
     */
    public VectorizeProgressInfo removeProgress(Long fileId) {
        return progressMap.remove(fileId);
    }

    /**
     * 获取所有文件的处理进度
     *
     * @return 所有进度信息对象的列表
     */
    public List<VectorizeProgressInfo> getAllProgress() {
        return new ArrayList<>(progressMap.values());
    }

    /**
     * 将Document对象转换为向量并存储到指定的向量存储中，并跟踪进度
     *
     * @param document       要处理的文档
     * @param embeddingStore 向量存储
     * @param embeddingModel 向量模型
     * @param fileId         文件ID，用于跟踪进度
     * @return 存储的向量数量
     */
    public int storeDocumentVectorsWithProgress(
            dev.langchain4j.data.document.Document document,
            dev.langchain4j.store.embedding.EmbeddingStore<dev.langchain4j.data.segment.TextSegment> embeddingStore,
            dev.langchain4j.model.embedding.EmbeddingModel embeddingModel,
            Long fileId) {

        try {
            String source = document.metadata().getString("source");
            log.info("开始向量化存储文档: {}, 文件ID: {}", source, fileId);

            // 更新进度状态为"处理中"
            updateProgress(fileId, 0, 0, "处理中");

            // 默认的分割大小和重叠
            int defaultMaxSegmentSize = 500;
            int defaultOverlap = 50;

            // 分割文档
            List<dev.langchain4j.data.segment.TextSegment> segments =
                    splitDocument(document, defaultMaxSegmentSize, defaultOverlap);

            // 更新总段落数
            updateProgress(fileId, 0, segments.size(), "开始向量化");

            // 向量化并存储段落
            int successCount = 0;
            for (dev.langchain4j.data.segment.TextSegment segment : segments) {
                try {
                    // 创建向量嵌入
                    dev.langchain4j.data.embedding.Embedding embedding = embeddingModel.embed(segment).content();

                    // 存储到向量数据库
                    embeddingStore.add(embedding, segment);
                    successCount++;

                    // 每处理10个段落或处理到最后一个，更新一次进度
                    if (successCount % 10 == 0 || successCount == segments.size()) {
                        updateProgress(fileId, successCount, segments.size(), "向量化处理中");
                    }
                } catch (Exception e) {
                    log.warn("段落向量化失败: {}", e.getMessage());
                }
            }

            // 标记处理完成
            completeProgress(fileId);

            log.info("文档向量化存储完成，成功存储 {} 个段落向量", successCount);
            return successCount;
        } catch (Exception e) {
            // 标记处理失败
            failProgress(fileId, e.getMessage());

            log.error("文档向量化存储失败: {}", e.getMessage(), e);
            throw new RuntimeException("文档向量化存储失败: " + e.getMessage(), e);
        }
    }
}
